package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"gitee.com/rocket049/baiduocr"
)

type Config struct {
	AppID     string
	ApiKey    string
	SecretKey string
}

const sampleJSON = `{
	"AppID":"xxxxxxx",
	"ApiKey":"xxxxxxxxxxxxx",
	"SecretKey":"xxxxxxxxxxxxxx"
}
`

func createSampleJSON(filename string) {
	fmt.Println("  配置文件：myocr.json，请填入你在百度大脑申请的文字识别应用的 APP_ID、API_KEY、SECRET_KEY\n", sampleJSON)
	ioutil.WriteFile(filename, []byte(sampleJSON), 0644)
}

func GetConfig() *Config {
	exe1, err := os.Executable()
	if err != nil {
		panic(err)
	}
	dir1 := filepath.Dir(exe1)
	cfgPath := filepath.Join(dir1, "myocr.json")
	cfgData, err := ioutil.ReadFile(cfgPath)
	if err != nil {
		createSampleJSON(cfgPath)
		panic(err)
	}
	res := new(Config)
	err = json.Unmarshal(cfgData, res)
	if err != nil {
		panic(err)
	}
	return res
}

func repeatOcr(client *baiduocr.Client, name string, writer io.Writer) {
	fmt.Println("准备好识别", name)
	fmt.Println("按回车键进行识别，输入‘q’退出")
	var s string
	for {
		fmt.Scanf("%s", &s)
		if s == "q" {
			break
		}
		getTextAndPrint(client, name, writer)
	}
}

func getTextAndPrint(client *baiduocr.Client, name string, writer io.Writer) {
	img, err := ioutil.ReadFile(name)
	if err != nil {
		panic(err)
	}
	res, err := client.General(img)
	if err != nil {
		panic(err)
	}
	var txt string
	if res.ErrorCode != 0 {
		//fmt.Println()
		txt = fmt.Sprintln(res.ErrorCode, res.ErrorMsg)
	} else {
		txt = fmt.Sprintf("%s\r\n", res.String())
	}
	writer.Write([]byte(txt))
}

func getJpgPngs() ([]string, error) {
	fp, err := os.Open(".")
	if err != nil {
		return nil, err
	}
	names, err := fp.Readdirnames(-1)
	if err != nil {
		return nil, err
	}
	res := []string{}
	for _, v := range names {
		vl := strings.ToLower(v)
		if strings.HasSuffix(vl, "jpg") || strings.HasSuffix(vl, "png") {
			res = append(res, v)
		}
	}
	sort.Strings(res)
	return res, nil
}

func myusage() {
	flag.Usage()
	fmt.Println("  myocr a.jpg b.png c.jpg d.png ...\n        连续识别各个文件，可以使用通配符，程序会按照文件名进行排序")
	fmt.Println("  myocr\n        连续识别当前目录中的 jpg、png 图片，程序会按照文件名进行排序")
	fmt.Println("  配置文件：myocr.json，请填入你在百度大脑申请的文字识别应用的 APP_ID、API_KEY、SECRET_KEY\n", sampleJSON)
}

func main() {
	var srcName = flag.String("src", "", "设置识别的图片文件名，用于反复识别同一个名字的图片")
	var fout = flag.String("o", "log.txt", "设置输出文件，保存识别结果")
	var h = flag.Bool("h", false, "显示帮助信息")
	flag.Parse()

	if *h {
		myusage()
		return
	}
	var writer io.Writer
	if len(*fout) > 0 {
		fp, err := os.Create(*fout)
		if err != nil {
			panic(err)
		}
		defer fp.Close()
		writer = io.MultiWriter(os.Stdout, fp)
	} else {
		writer = os.Stdout
	}

	cfg := GetConfig()
	client, err := baiduocr.NewClient(cfg.AppID, cfg.ApiKey, cfg.SecretKey)
	if err != nil {
		fmt.Println("  配置文件：myocr.json，请填入你在百度大脑申请的文字识别应用的 APP_ID、API_KEY、SECRET_KEY\n", sampleJSON)
		panic(err)
	}

	if *srcName != "" {
		repeatOcr(client, *srcName, writer)

	} else if len(flag.Args()) == 1 {
		getTextAndPrint(client, flag.Arg(0), writer)

	} else if len(flag.Args()) > 1 {
		files := flag.Args()
		sort.Strings(files)
		for _, v := range files {
			getTextAndPrint(client, v, writer)
			time.Sleep(time.Second / 2)
		}

	} else {
		names, err := getJpgPngs()
		if err != nil {
			panic(err)
		}
		for _, v := range names {
			getTextAndPrint(client, v, writer)
			time.Sleep(time.Second / 2)
		}
	}
}
